using Unity.Burst;                     // 引入 Burst 编译器相关功能，用于高性能并行计算
using Unity.Collections;               // 提供 Native 容器（如 NativeArray），用于高效内存管理
using Unity.Jobs;                      // 提供 Job 系统，用于多线程任务调度
using Unity.Mathematics;               // 提供数学库，如 float3, int4 等类型及数学函数
using UnityEngine;                     // Unity 主命名空间

using static Unity.Mathematics.math;   // 静态导入 math 命名空间，方便直接使用 math.xxx 函数

// 该类用于可视化基于哈希值的实例化对象分布，比如在空间中按照某种规则散布物体并赋予其唯一性
public class HashVisualization : MonoBehaviour
{
	/// <summary>
	/// 一个使用 Burst 编译的 Job，用于并行计算一组变换后位置的哈希值。
	/// 每个 Job 实例处理一部分数据（IJobFor 支持并行 for 循环）。
	/// </summary>
	[BurstCompile(FloatPrecision.Standard, FloatMode.Fast, CompileSynchronously = true)]
	struct HashJob : IJobFor
	{

		/// <summary>
		/// 只读：存储了一组 4x3 的位置矩阵（每行是一个 float3，共 4 个点，即一个 float3x4）。
		/// </summary>
		[ReadOnly]
		public NativeArray<float3x4> positions;

		/// <summary>
		/// 可写：用于存储计算出来的哈希值，每个 uint4 存储 4 个 uint 类型的哈希。
		/// </summary>
		[WriteOnly]
		public NativeArray<uint4> hashes;

		/// <summary>
		/// 用于生成哈希值的小型 XXHash 算法工具，支持一次处理多个值（uint4）。
		/// </summary>
		public SmallXXHash4 hash;

		/// <summary>
		/// 定义了一个变换矩阵（float3x4 格式），用于将原始位置从一个空间转换到另一个空间（比如局部到世界）。
		/// </summary>
		public float3x4 domainTRS;

		/// <summary>
		/// 将一个 float4x3 的位置矩阵，根据给定的 float3x4 变换矩阵（TRS）进行变换。
		/// 返回一个新的 float4x3，表示变换后的位置。
		/// </summary>
		/// <param name="trs">变换矩阵，格式为 float3x4</param>
		/// <param name="p">输入的位置矩阵，格式为 float4x3（4个点，每个点3个坐标）</param>
		/// <returns>变换后的位置矩阵，仍然是 float4x3</returns>
		float4x3 TransformPositions(float3x4 trs, float4x3 p) =>
			float4x3(
				trs.c0.x * p.c0 + trs.c1.x * p.c1 + trs.c2.x * p.c2 + trs.c3.x,  // X 轴变换
				trs.c0.y * p.c0 + trs.c1.y * p.c1 + trs.c2.y * p.c2 + trs.c3.y,  // Y 轴变换
				trs.c0.z * p.c0 + trs.c1.z * p.c1 + trs.c2.z * p.c2 + trs.c3.z   // Z 轴变换
			);

		/// <summary>
		/// 对索引为 i 的一组 4 个点，计算它们的哈希值
		/// </summary>
		/// <param name="i">当前处理的索引，对应 positions 和 hashes 数组中的第 i 项</param>
		public void Execute(int i)
		{
			// 先将 positions[i]（float3x4）转置成 float4x3，然后应用域变换
			float4x3 p = TransformPositions(domainTRS, transpose(positions[i]));

			// 对每个点的 X、Y、Z 分量分别取 floor 后转换为 int，得到整数网格坐标 u, v, w
			int4 u = (int4)floor(p.c0);  // X 坐标 -> 网格 X
			int4 v = (int4)floor(p.c1);  // Y 坐标 -> 网格 Y
			int4 w = (int4)floor(p.c2);  // Z 坐标 -> 网格 Z

			// 使用 SmallXXHash4 工具，依次“吃下” u, v, w 整数坐标，生成一组 4 个哈希值
			hashes[i] = hash.Eat(u).Eat(v).Eat(w);
		}
	}

	public enum Shape { Plane, Sphere, Torus }

	// ==================== 静态字段：形状 Job 调度委托数组 ====================

	/// <summary>
	/// 一个静态数组，存储了三种不同形状（平面、球体、环面）对应的 Job 调度委托。
	/// 每个委托用于在并行中生成相应形状上的点位置和法线。
	/// </summary>
	static Shapes.ScheduleDelegate[] shapeJobs = {
		Shapes.Job<Shapes.Plane>.ScheduleParallel,   // 平面形状的 Job 调度
        Shapes.Job<Shapes.Sphere>.ScheduleParallel,  // 球体形状的 Job 调度
        Shapes.Job<Shapes.Torus>.ScheduleParallel    // 环面形状的 Job 调度
    };

	// ==================== Shader 相关常量 ID ====================

	/// <summary>
	/// 在 Shader 中用于访问 _Hashes 缓冲区的属性 ID。
	/// </summary>
	static int
		hashesId = Shader.PropertyToID("_Hashes"),
		/// <summary>
		/// 在 Shader 中用于访问 _Positions 缓冲区的属性 ID。
		/// </summary>
		positionsId = Shader.PropertyToID("_Positions"),
		/// <summary>
		/// 在 Shader 中用于访问 _Normals 缓冲区的属性 ID。
		/// </summary>
		normalsId = Shader.PropertyToID("_Normals"),
		/// <summary>
		/// 在 Shader 中用于访问 _Config 向量的属性 ID，通常包含分辨率、缩放、位移等配置信息。
		/// </summary>
		configId = Shader.PropertyToID("_Config");

	// ==================== Unity Inspector 可配置字段 ====================

	/// <summary>
	/// 在场景中显示的实例化网格预制体。
	/// </summary>
	[SerializeField]
	Mesh instanceMesh;

	/// <summary>
	/// 应用于实例化对象的材质。
	/// </summary>
	[SerializeField]
	Material material;

	/// <summary>
	/// 当前选择的几何形状类型：平面、球体、环面。
	/// </summary>
	[SerializeField]
	Shape shape;

	/// <summary>
	/// 实例化对象相对于网格大小的缩放比例。
	/// </summary>
	[SerializeField, Range(0.1f, 10f)]
	float instanceScale = 2f;

	/// <summary>
	/// 空间分辨率，即每行/列生成的实例数量，最终实例总数为 resolution * resolution。
	/// </summary>
	[SerializeField, Range(1, 512)]
	int resolution = 16;

	/// <summary>
	/// 用于给实例添加随机扰动的位移强度。
	/// </summary>
	[SerializeField, Range(-0.5f, 0.5f)]
	float displacement = 0.1f;

	/// <summary>
	/// 用于初始化随机哈希种子的值。
	/// </summary>
	[SerializeField]
	int seed;

	/// <summary>
	/// 定义了一个变换空间（TRS：平移、旋转、缩放），用来控制实例分布的空间范围和变换。
	/// </summary>
	[SerializeField]
	SpaceTRS domain = new SpaceTRS
	{
		scale = 8f  // 默认空间缩放为 8 单位
	};

	// ==================== Native 数据容器 ====================

	/// <summary>
	/// 存储每个实例位置对应的哈希值，每个元素是 uint4，即 4 个哈希。
	/// </summary>
	NativeArray<uint4> hashes;

	/// <summary>
	/// 存储每个实例的位置信息，每个元素是 float3x4，即 4 个点的位置。
	/// </summary>
	NativeArray<float3x4> positions, normals;

	// ==================== ComputeBuffer 缓冲区 ====================

	/// <summary>
	/// 用于将哈希数据传递给 GPU Shader 的 ComputeBuffer。
	/// </summary>
	ComputeBuffer hashesBuffer, positionsBuffer, normalsBuffer;

	// ==================== MaterialPropertyBlock 与 Bounds ====================

	/// <summary>
	/// 用于在不创建新材质实例的情况下，动态设置材质参数。
	/// </summary>
	MaterialPropertyBlock propertyBlock;

	/// <summary>
	/// 标记当前是否需要重新生成数据（比如参数改变或物体移动）。
	/// </summary>
	bool isDirty;

	/// <summary>
	/// 定义了所有实例化对象的包围盒，用于优化渲染剔除。
	/// </summary>
	Bounds bounds;

	// ==================== 生命周期方法 ====================

	/// <summary>
	/// 当脚本启用时调用，初始化所有 Native 数组和 ComputeBuffer。
	/// </summary>
	void OnEnable()
	{
		isDirty = true;  // 标记需要更新

		// 计算需要的长度：总实例数 = resolution * resolution，然后按 4 个一组分组，向上取整
		int length = resolution * resolution;
		length = length / 4 + (length & 1);  // 如果不能被4整除，则多分配一组

		// 初始化 Native 数组，用于存储哈希、位置、法线
		hashes = new NativeArray<uint4>(length, Allocator.Persistent);
		positions = new NativeArray<float3x4>(length, Allocator.Persistent);
		normals = new NativeArray<float3x4>(length, Allocator.Persistent);

		// 初始化 ComputeBuffer，用于将数据传入 GPU
		hashesBuffer = new ComputeBuffer(length * 4, 4);           // 每个 uint4 是 4 * 4=16 字节？或者每个 uint 是 4 字节，4个就是16字节？
		positionsBuffer = new ComputeBuffer(length * 4, 3 * 4);    // 每个 float3 是 3 * 4=12 字节，4个就是48字节
		normalsBuffer = new ComputeBuffer(length * 4, 3 * 4);

		// 初始化 MaterialPropertyBlock，用于设置 Shader 参数
		propertyBlock ??= new MaterialPropertyBlock();
		propertyBlock.SetBuffer(hashesId, hashesBuffer);
		propertyBlock.SetBuffer(positionsId, positionsBuffer);
		propertyBlock.SetBuffer(normalsId, normalsBuffer);

		// 设置一个配置向量，包括分辨率、缩放比例、位移
		propertyBlock.SetVector(configId, new Vector4(
			resolution, instanceScale / resolution, displacement
		));
	}

	/// <summary>
	/// 当脚本禁用时调用，释放所有 Native 和 ComputeBuffer 资源。
	/// </summary>
	void OnDisable()
	{
		hashes.Dispose();
		positions.Dispose();
		normals.Dispose();
		hashesBuffer.Release();
		positionsBuffer.Release();
		normalsBuffer.Release();

		hashesBuffer = null;
		positionsBuffer = null;
		normalsBuffer = null;
	}

	/// <summary>
	/// 当在 Inspector 中修改了参数时调用，如果当前启用了，则重新初始化资源。
	/// </summary>
	void OnValidate()
	{
		if (hashesBuffer != null && enabled)
		{
			OnDisable();
			OnEnable();
		}
	}

	/// <summary>
	/// 每帧调用，如果数据有变动或物体发生了移动，则重新生成哈希和位置数据，并进行绘制。
	/// </summary>
	void Update()
	{
		if (isDirty || transform.hasChanged)
		{
			isDirty = false;
			transform.hasChanged = false;

			// 调用对应的 ShapeJob，生成 positions 和 normals 数据
			JobHandle handle = shapeJobs[(int)shape](
				positions, normals, resolution, transform.localToWorldMatrix, default
			);

			// 并行计算哈希值：对每个 float3x4（4个点）计算一个 uint4 哈希
			new HashJob
			{
				positions = positions,
				hashes = hashes,
				hash = SmallXXHash.Seed(seed),  // 用种子初始化哈希生成器
				domainTRS = domain.Matrix       // 应用域变换矩阵
			}.ScheduleParallel(hashes.Length, resolution, handle).Complete();  // 并行调度并等待完成

			// 将 Native 数据拷贝到 ComputeBuffer，以供 GPU 使用
			hashesBuffer.SetData(hashes.Reinterpret<uint>(4 * 4));          // 每个 uint4 是 4个uint，每个4字节 => 16字节？可能需确认
			positionsBuffer.SetData(positions.Reinterpret<float3>(3 * 4 * 4));
			normalsBuffer.SetData(normals.Reinterpret<float3>(3 * 4 * 4));

			// 动态计算包围盒，用于渲染优化
			bounds = new Bounds(
				transform.position,
				float3(2f * cmax(abs(transform.lossyScale)) + displacement)
			);
		}

		// 使用 Graphics.DrawMeshInstancedProcedural 绘制所有实例化网格
		Graphics.DrawMeshInstancedProcedural(
			instanceMesh, 0, material, bounds, resolution * resolution, propertyBlock
		);
	}
}